# llms/database.py 数据库连接和配置
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker, scoped_session
from sqlalchemy.ext.declarative import declarative_base
from config.db_config import Config

# 数据库连接配置
SQLALCHEMY_DATABASE_URL = Config.SQLALCHEMY_DATABASE_URI
# 创建数据库引擎
engine = create_engine(SQLALCHEMY_DATABASE_URL, echo=True, pool_pre_ping=True)

# 创建数据库会话工厂
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

# 创建数据库基类
Base = declarative_base()

# 创建数据库会话
db_session = scoped_session(SessionLocal)


def get_db():
    """在FastAPI中使用依赖注入时，该函数会被调用以获取数据库会话"""
    db = db_session
    try:
        yield db
    finally:
        db.close()


# 测试数据库连接
def test_db_connection():
    try:
        db_session.execute(text("SELECT 1"))
        print("数据库连接正常")
    except Exception as e:
        print(f"数据库连接异常: {e}")
